package main

import (
	"bytes"
	"fmt"
	"log"
	"os/exec"
	"strconv"
	"strings"
)

func runCommand(cmd string) (string, error) {
	var out bytes.Buffer
	var stderr bytes.Buffer
	cmdObj := exec.Command("bash", "-c", cmd)
	cmdObj.Stdout = &out
	cmdObj.Stderr = &stderr
	err := cmdObj.Run()
	if err != nil {
		log.Println("命令执行出错: ", cmd, err)
		return "", err
	}
	return out.String(), nil
}

func getSpecificGpuCacheUtilization(nf string) float64 {
	nvidiaOutput, err := runCommand(fmt.Sprintf("nvidia-smi --query-compute-apps=pid,used_memory --format=csv,noheader,nounits"))
	if err != nil {
		return 0
	}
	nvidiaLines := strings.Split(nvidiaOutput, "\n")
	
	gpumem := map[string]float64{}
	for _, line := range nvidiaLines {
		ss := strings.Split(line, ", ")
		if len(ss) < 2 {
			continue
		}
		gpumem[ss[0]], _ = strconv.ParseFloat(ss[1], 64)
	}
	

	// get pids
	pidsOutput, err := runCommand(fmt.Sprintf("pstree -p  $(docker inspect -f '{{.State.Pid}}' $(docker ps -a | grep %s | awk '{print $1}')) | grep -oP '\\(\\K[0-9]+'", nf))
	if err != nil {
		return 0
	}
	pids := strings.Split(pidsOutput, "\n")
	// fmt.Println(gpumem, pids)
	for _, spid := range pids {
		if v, ok := gpumem[spid]; ok {
			return v
		}
	}
	return 0
}

func main() {
	fmt.Println(getSpecificGpuCacheUtilization("ffmpeg-container_ffmpeg-gpu-pod-1"))
	fmt.Println(getSpecificGpuCacheUtilization("upf_upf-0"))
	fmt.Println(getSpecificGpuCacheUtilization("upf_upf-1"))
}
